import pymysql
from pymysql import Error

def create_connection():
    """创建数据库连接"""
    try:
        connection = pymysql.connect(
            host='127.0.0.1',      # 数据库主机地址
            user='root',  # 数据库用户名
            port=3306,  #数据库服务器端口
            password='root',  # 数据库密码
            database='aichat',  # 数据库名称
            charset='utf8mb4',     # 字符编码
            cursorclass=pymysql.cursors.DictCursor  # 返回字典格式的结果
        )
        print("成功连接到MySQL数据库")
        return connection
    except Error as e:
        print(f"连接错误: {e}")
        return None

def insert_chat_content(identity, content, theme_id, theme_name):
    """插入聊天内容到数据库
    identity: 'user' 或 'ai'
    content: 聊天内容
    theme_id: 主题ID
    theme_name: 主题名称
    """
    connection = None
    cursor = None
    try:
        connection = create_connection()
        if connection is None:
            raise Error("数据库连接失败")
        
        cursor = connection.cursor()
        sql = 'INSERT INTO chat_content (identity, content, theme_id, theme_name) VALUES (%s, %s, %s, %s)'
        cursor.execute(sql, (identity, content, theme_id, theme_name))
        connection.commit()
        print(f"成功插入{identity}的聊天内容")
        return True
    except Error as e:
        print(f"插入聊天内容错误: {e}")
        if connection and connection.open:
            connection.rollback()
        return False
    finally:
        if cursor:
            cursor.close()
        if connection and connection.open:
            connection.close()

# 使用示例
if __name__ == "__main__":
    # 主题名称
    theme_name = "金融服务"
    connection = None
    cursor = None
    theme_id = 0
    
    try:
        connection = create_connection()
        if connection is None:
            raise Error("数据库连接失败")
        
        cursor = connection.cursor()

        # 1. 检查主题是否存在
        cursor.execute("SELECT id FROM chat_themes WHERE theme_name = %s", (theme_name,))
        result = cursor.fetchone()

        if result:
            # 主题存在，获取id
            theme_id = result['id']
            print(f"主题 '{theme_name}' 已存在，ID: {theme_id}")
        else:
            # 主题不存在，创建新主题
            cursor.execute("INSERT INTO chat_themes (theme_name) VALUES (%s)", (theme_name,))
            connection.commit()
            theme_id = cursor.lastrowid
            print(f"创建新主题 '{theme_name}'，ID: {theme_id}")
    except Error as e:
        print(f"主题处理错误: {e}")
    finally:
        if cursor:
            cursor.close()
        if connection and connection.open:
            connection.close()

    # 2. 插入用户提问
    insert_chat_content("user", "123123", theme_id, theme_name)
    
    # 3. 插入AI回答
    insert_chat_content("ai", "这是AI的回答内容", theme_id, theme_name)

